refactor(attn): let the bucketer own the batched decode buckets - #810

👋 Hi! Thank you for contributing.
Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, run ./format.sh.
Now you are good to go 🚀.

We also recommend installing prek and configuring it to check your code before every local commit.

I also ran it through Claude an it additionally found those points:

1. Defaults do change for non-power-of-two block_size (highest priority)
   The review confirms defaults are unchanged; that holds only for power-of-two block_size. At block_size=192 — supported, warned about in spyre_attn_bucketer.py:131-140, and tested in tests/runtime/test_spyre_attn_bucketer.py:77-82 — the ladder moves with no env override set:

mml=2048 [1, 2, 4, 8, 11] -> [2, 3, 6, 11]
mml=32768 [1,2,4,8,16,32,64,128,171] -> [2,3,6,11,22,43,86,171]
mml=256 [1, 2] -> [2]
The floor rises 1 → 2 and the ladder gets sparser at the top. This reframes the spyre_attn.py:645 thread from "an opt-in env var moves this" to "a supported default config regresses", which argues for keeping the decode ladder independent rather than documenting the trade-off. No test catches it because nothing asserts num_blocks_buckets under a non-power-of-two block_size; the -k batched_decode run uses block_size=128, where the ladder is genuinely identical.

2. Overrides above max_model_len are pathological, and a test sanctions them
   SPYRE_ATTN_KV_BUCKETS=128,8192 at max_model_len=2048, block_size=64 gives num_blocks_buckets == [2, 128] (was [1,2,4,8,16,32]). A 3-block decode then runs 128 kernel iterations — ~40× worse again than the 32× case on the :645 thread. tests/runtime/test_spyre_attn_bucketer.py:242-247 deliberately preserves above-limit entries, so this is a reachable, blessed configuration.

3. The per-step H2D cost scales too, not just gathers and matmuls
   mask_by_block_cpu is sized [b_blocks, b_seqs*KV, 1, block_size] (spyre_attn.py:1068-1090) and mirrored to device on every step. Raising b_blocks inflates that transfer by the same factor as the loop count, so the :645 cost estimate is a lower bound.

4. Correctness is preserved — recording why, so this isn't read as a soundness bug
   n_use = min(n, b_blocks) (spyre_attn.py:1054-1055) bounds the fill by real blocks, so block_table is never over-indexed; every row keeps a finite block 0 (real rows by construction, padded rows via mask_bs_bb[num_seqs:, 0] = finfo.min at spyre_attn.py:1081), so exp(-inf - finite) == 0 and new_max is never -inf - -inf. A larger b_blocks only adds zero-contribution blocks. Everything above is throughput and transfer volume, not wrong output.

5. The "fewer compiles" half of the corrected benefit is zero by default
   In the default config the ladder is byte-identical, so there are exactly as many distinct b_blocks values as before — no compile reduction. It becomes nonzero only where the ladder shortens, and it shortens by raising the floor (512,1024,2048: 6 block counts → 3, floor 1 → 8). The claimed benefit and the regression are the same phenomenon. For the default path the defensible claim is the single source of truth on its own — which is reason enough to land the refactor, and an argument for landing it without coupling the decode ladder to the KV axis.

6. The new properties hand out mutable internal state
   num_seqs_buckets / num_blocks_buckets (spyre_attn_bucketer.py:200-206) return the live list. Making them tuple[int, ...] also touches the == comparison at spyre_model_runner.py:797.

7. _resolve_builder_attn_bucketer doesn't compare the new axis
   spyre_model_runner.py:797-808 asserts (block_size, num_blocks_buckets) agree across builders. Now that the builder's num_seqs ladder comes from the bucketer too, either add it to the tuple or say in the docstring why it's exempt (nothing records that axis).

8. A third copy of the same ladder still exists
   spyre_shape_bucketer.batch_buckets (spyre_shape_bucketer.py:94-108) reimplements it, and its docstring at line 97 cross-references _powers_of_two_up_to — a pointer this PR silently redirects to another module. They also diverge at the edge: batch_buckets(0) == [1], _powers_of_two_up_to(0) == (). In scope for a "single source of truth" PR, or at least a TODO naming the new location.

9. Dropped TODO, and num_seqs is now the only axis without an env knob
   The deleted comment carried TODO: expose as engine args if configurability is needed. KV and query both route through _resolve_buckets with an override; num_seqs silently doesn't. Either add one for symmetry or carry the TODO forward.

10. The _powers_of_two_up_to docstring lost two sentences in the move
    spyre_attn_bucketer.py:72-84 dropped "plus n itself if it is not already a power of 2" and "A start above n yields just (n,)." The second documents exactly the edge case that the start=_MIN_BATCHED_SEQS suggestion creates.

11. Module and class docstrings now overclaim
    spyre_attn_bucketer.py:15-29 says the module enumerates pairs "so warmup can record them all up front", and the class docstring says it enumerates "the variants to record". num_seqs_buckets is dispatch-only and never recorded. Worth stating at the new field, in the class docstring, and in variants() (why the axis is deliberately absent).

12. SpyreAttnBucket docstring is now false for one kernel
    spyre_attn_bucketer.py:50-59: "Fields are the values the kernel specializes on, so a recorded bucket and a runtime dispatch reach the same Dynamo entry." True for _attn_fn, false for _batched_decode_kernel, which specializes on (num_seqs, num_blocks).

13. The new builder lines are uncommented
    spyre_attn.py:644-645 — the deleted lines explained what these were ("Buckets for the batched decode fast path. One compiled kernel per bucket."); the replacements explain nothing, and they sit under a comment block about _attn_bucketer. Moot if the :644 thread's suggestion lands.

⚙️ Auto-merge enabled — dispatched /spyre-test for 68102bc9 because the required Spyre Test gate had not run on this commit.

_{Automated: enabling auto-merge requests the gate, so the PR is not left waiting on a check nobody asked for. Re-enabling auto-merge on this same commit will not dispatch again.}

🔄 spyre-test: running

Triggered by: /spyre-test comment

flowchart LR
  subgraph Lamd64_0["amd64 L0 · 1 parallel"]
    n_amd64_flex["flex ↺<br/>rpm · 21f9dbf1c01c"]
  end
  subgraph Lamd64_1["amd64 L1 · 3 parallel"]
    n_amd64_aiu_toolbox_ibm_aiu_toolbox_e2e["aiu-toolbox/ibm-aiu-toolbox-e2e<br/>rpm · 6068750b7f59"]
    n_amd64_libaiupti["libaiupti<br/>rpm · 47e615e55dc3"]
    n_amd64_spyre_comms["spyre-comms<br/>rpm · cd26cd09b611"]
  end
  subgraph Lamd64_2["amd64 L2 · 1 parallel"]
    n_amd64_spyre_backend_spyre_backend_dev["spyre-backend/spyre-backend-dev<br/>image · 00cdede56c88"]
  end
  subgraph Lamd64_3["amd64 L3 · 1 parallel"]
    n_amd64_torch_spyre_torch_spyre_dev["torch-spyre/torch-spyre-dev<br/>image · dd7d61b76437"]
  end
  subgraph Lamd64_4["amd64 L4 · 1 parallel"]
    n_amd64_hf_adapters_hf_adapters_dev["hf-adapters/hf-adapters-dev<br/>image · fa2a54db6cab"]
  end
  subgraph Lamd64_5["amd64 L5 · 1 parallel"]
    n_amd64_spyre_inference_spyre_inference_dev["spyre-inference/spyre-inference-dev<br/>image · d98eba460635"]
  end
  n_amd64_flex --> n_amd64_spyre_comms
  n_amd64_flex --> n_amd64_aiu_toolbox_ibm_aiu_toolbox_e2e
  n_amd64_flex --> n_amd64_libaiupti
  n_amd64_aiu_toolbox_ibm_aiu_toolbox_e2e --> n_amd64_spyre_backend_spyre_backend_dev
  n_amd64_flex --> n_amd64_spyre_backend_spyre_backend_dev
  n_amd64_libaiupti --> n_amd64_spyre_backend_spyre_backend_dev
  n_amd64_spyre_comms --> n_amd64_spyre_backend_spyre_backend_dev
  n_amd64_spyre_backend_spyre_backend_dev --> n_amd64_torch_spyre_torch_spyre_dev
  n_amd64_torch_spyre_torch_spyre_dev --> n_amd64_hf_adapters_hf_adapters_dev
  n_amd64_hf_adapters_hf_adapters_dev --> n_amd64_spyre_inference_spyre_inference_dev
  classDef sPending fill:#eceff1,stroke:#90a4ae,color:#37474f
  classDef sBuilding fill:#fff8e1,stroke:#f9a825,color:#5d4037,stroke-width:2px
  classDef sOk fill:#e8f5e9,stroke:#43a047,color:#1b5e20
  classDef sReused fill:#e3f2fd,stroke:#1e88e5,color:#0d47a1
  classDef sFailed fill:#ffebee,stroke:#e53935,color:#b71c1c,stroke-width:2px
  classDef sDropped fill:#f5f5f5,stroke:#bdbdbd,color:#9e9e9e
  class n_amd64_flex sReused;
  class n_amd64_aiu_toolbox_ibm_aiu_toolbox_e2e sBuilding;
  class n_amd64_libaiupti sBuilding;
  class n_amd64_spyre_comms sBuilding;
  class n_amd64_spyre_backend_spyre_backend_dev sPending;
  class n_amd64_torch_spyre_torch_spyre_dev sPending;
  class n_amd64_hf_adapters_hf_adapters_dev sPending;
  class n_amd64_spyre_inference_spyre_inference_dev sPending;

🔨 orch trigger-pr-validation · arches amd64 · fp amd64=e261d97c